package com.aispace.taichu.modules.agent.llm;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.aispace.taichu.modules.dynamic.model.model.entity.ModelConfig;
import com.aispace.taichu.modules.dynamic.model.service.IModelConfigService;
import com.aispace.taichu.modules.listener.ModelChangeEvent;
import com.aispace.taichu.modules.listener.TaiChuEvent;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@Slf4j
@RequiredArgsConstructor
public class LLMService implements ILLmService{

    private final IModelConfigService modelService;

    private ChatClient chatClient;

    private Map<Long, ChatClient> clients = new ConcurrentHashMap<>();

    @Override
    public ChatClient getDefaultChatClient() {
        QueryWrapper<ModelConfig> wrapper = new QueryWrapper<>();
        wrapper.lambda().eq(ModelConfig::getIsDefault, true);
        ModelConfig defaultModel = modelService.getOne(wrapper);
        if (defaultModel == null) {
            log.error("No default model found");
            List<ModelConfig> models = modelService.list();
            if (models.isEmpty()) {
                log.error("No model found");
            } else {
                defaultModel = models.stream().findFirst().orElse(null);
                if (defaultModel != null) {
                    log.warn("Using first model as default: {}", defaultModel.getModelName());
                } else {
                    log.error("First model in list is also null");
                }
            }
        }else {
            log.info("Find default model: {}", defaultModel.getModelName());
        }
        if (defaultModel != null) {
           return getChatClient(defaultModel);
        }
        else {
            log.warn("Cannot find any model，ChatClient will be initialize after model being configured");
            return null;
        }
    }

    @Override
    public ChatClient getChatClient(ModelConfig  model) {
        Long modelId = model.getId();
        if (clients.containsKey(modelId)) {
            return clients.get(modelId);
        }
        return this.buildOrUpdateChatClient(model);
    }

    @PostConstruct
    public void initializeChatClients() {
        ChatClient client = getDefaultChatClient();
        if (client != null) {
            this.chatClient = client;
            log.info("Initialize default chat client success");
        }
        else {
            log.warn("Cannot find any model，ChatClient will be initialize after model being configured");
        }
    }

    private void initializeChatClients(ModelConfig model) {
        this.chatClient=buildOrUpdateChatClient(model);
    }

    private ChatClient buildOrUpdateChatClient(ModelConfig model) {
        Map<String, String> headers = model.getRequestHeaders();
        OpenAiApi openAiApi = OpenAiApi.builder()
                .baseUrl(model.getBaseUrl())
                .apiKey(model.getApiKey())
                .completionsPath(model.getCompletionsPath())
                .build();
        OpenAiChatOptions.Builder chatOptionsBuilder = OpenAiChatOptions.builder().model(model.getModelName());
        if (model.getTemperature() != null) {
            chatOptionsBuilder.temperature(model.getTemperature());
        }
        if (model.getTopP() != null) {
            chatOptionsBuilder.topP(model.getTopP());
        }
        chatOptionsBuilder.internalToolExecutionEnabled(false);
        OpenAiChatOptions chatOptions = chatOptionsBuilder.build();
        if (headers != null) {
            chatOptions.setHttpHeaders(headers);
        }
        OpenAiChatModel openAiChatModel = OpenAiChatModel.builder()
                .openAiApi(openAiApi)
                .defaultOptions(chatOptions)
                .build();
        ChatClient client = ChatClient.builder(openAiChatModel)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();
        clients.put(model.getId(), client);
        log.info("Build or update dynamic chat client for model: {}", model.getModelName());
        return client;
    }

    @EventListener
    public void onEvent(TaiChuEvent<ModelChangeEvent> event) {
        log.info("Received ModelChangeEvent...");
        if (event.getEntity() == null) {
            log.warn("ModelChangeEvent entity is null");
            return;
        }
        log.info("Model updated");
        ModelConfig modelEntity = event.getEntity().getModelConfig();
        initializeChatClients(modelEntity);
    }

}
